import asyncio
import uuid
import os
import pickle
import threading
from functools import partial
from fastapi import FastAPI, WebSocket, WebSocketDisconnect
from typing import List ,Dict,Optional
from bisheng.cache.redis import redis_client
from bisheng.utils.logger import logger



# class WebsocketManager:
#     def __init__(self):
#         self.local_connections: Dict[str, WebSocket] = {}  # 本地连接池
#         self.node_id = str(uuid.uuid4())  # 当前节点唯一标识
#         self.channel = "websocket_broadcast"  # Redis 频道名称
#
#         # 订阅 Redis 频道（重要！）
#         logger.info(f"-------订阅redis频道Connected: {self.node_id}-----------")
#         redis_client.subscribe(self.channel, self._handle_redis_message)
#
#     async def connect(self, websocket: WebSocket):
#         """处理 WebSocket 连接"""
#         await websocket.accept()
#         ws_id = str(uuid.uuid4())
#         self.local_connections[ws_id] = websocket
#         logger.info(list(self.local_connections.values()))
#         logger.info(f"Connected: {ws_id}")
#
#     def disconnect(self, websocket: WebSocket):
#         """处理 WebSocket 断开"""
#         for ws_id, ws in list(self.local_connections.items()):
#             if ws == websocket:
#                 del self.local_connections[ws_id]
#                 logger.info(f"Disconnected: {ws_id}")
#                 break
#
#     def broadcast(self, message: str):
#         """广播消息（入口方法）"""
#         # 1. 发送本地连接
#         # await self._send_to_local(message)
#
#         # 2. 通知其他节点
#         self._notify_other_nodes(message)
#
#     def _send_to_local(self, message: str):
#         logger.info(f"-------_send_to_local: {message}---------")
#         logger.info(f"----------{self.local_connections.items()}----------")
#         """发送给本地所有连接"""
#         for ws_id, ws in list(self.local_connections.items()):
#             try:
#                 self.send_personal_message(message,ws)
#             except Exception as e:
#                 logger.error(f"Send failed: {str(e)}")
#                 self.disconnect(ws)
#
#
#     def _notify_other_nodes(self, message: str):
#         """通过 Redis 发布消息"""
#         payload = {
#             "source": self.node_id,
#             "message": message,
#             "timestamp": asyncio.get_event_loop().time()
#         }
#         redis_client.publish(self.channel, pickle.dumps(payload))
#
#     def _handle_redis_message(self, channel: bytes, data: bytes, _):
#         """处理来自 Redis 的订阅消息（回调函数）"""
#         # 确保频道匹配且消息有效
#         # if channel.decode() != self.channel or not data:
#         #     return
#         try:
#
#             payload = pickle.loads(data)
#             logger.info(f"Redis message received: {payload}")
#             logger.info(f"-----------------redis 处理消息---------------")
#             # logger.info(f"主循环状态: Running={current_loop.is_running()}, Closed={current_loop.is_closed()}")
#             # logger.info(f"主循环状态: Running={current_loop1.is_running()}, Closed={current_loop1.is_closed()}")
#
#             #忽略自己发出的消息（避免循环）
#             # if payload["source"] == self.node_id:
#             #     return
#
#             self._send_to_local(payload["message"])
#             # 将同步方法转为异步执行
#             # future = asyncio.run_coroutine_threadsafe(self._send_to_local(payload["message"]),current_loop1)
#             # # 添加完成回调记录日志
#             # future.add_done_callback(lambda f: logger.info("任务执行完成，异常: %s", f.exception()))
#         except Exception as e:
#             logger.error(f"处理 Redis 消息失败: {str(e)}")
#
#     @staticmethod
#     def send_personal_message(message: str, websocket: WebSocket):
#         if websocket.client_state.CONNECTED:  # 检查连接是否仍然打开
#             websocket.send_text(message)
#
# # 创建全局单例实例
# ws_manager = WebsocketManager()


class WebsocketManager:
    def __init__(self):
        self.local_connections: Dict[str, WebSocket] = {}
        self.node_id = str(uuid.uuid4())
        self.channel = "websocket_broadcast"
        # 确保使用正确的事件循环
        self.loop = asyncio.get_event_loop()
        redis_client.subscribe(self.channel, self._handle_redis_message)
        logger.info(f"-------订阅redis频道Connected: {self.node_id}-----------")
    async def connect(self, websocket: WebSocket):
        await websocket.accept()
        ws_id = str(uuid.uuid4())
        self.local_connections[ws_id] = websocket
        logger.info(f"Connected: {ws_id}, 当前连接数: {len(self.local_connections)}")
    def disconnect(self, websocket: WebSocket):
        for ws_id, ws in list(self.local_connections.items()):
            if ws == websocket:
                del self.local_connections[ws_id]
                logger.info(f"Disconnected: {ws_id}")
                break
    def broadcast(self, message: str):
        # 发送到本地和其他节点
        self._send_to_local(message)
        self._notify_other_nodes(message)
    def _send_to_local(self, message: str):
        logger.info(f"尝试发送消息到本地连接，当前连接: {self.local_connections}")
        logger.info(f"------发送消息，{message}")
        for ws_id, ws in list(self.local_connections.items()):
            try:
                # await self.send_personal_message(message, ws)
                asyncio.run_coroutine_threadsafe(
                    self.send_personal_message(message, ws),  # 异步方法
                    self.loop  # 确保 loop 是主事件循环
                )
            except Exception as e:
                logger.error(f"发送失败: {e}")
                self.disconnect(ws)
    def _notify_other_nodes(self, message: str):
        payload = {
            "source": self.node_id,
            "message": message,
            "timestamp": asyncio.get_event_loop().time()
        }
        redis_client.publish(self.channel, pickle.dumps(payload))
    def _handle_redis_message(self, channel: bytes, data: bytes, pattern: Optional[bytes] = None):
        try:
            payload = pickle.loads(data)
            str_channel = channel.decode('utf-8')
            if str_channel == self.channel:
                logger.info(f"Connected: {self.local_connections},当前连接数: {len(self.local_connections)}")
                if payload["source"] == self.node_id:
                    return  # 避免循环
                self._send_to_local(payload["message"])
                # 调度到主事件循环执行
                # asyncio.run_coroutine_threadsafe(
                #     self._send_to_local(payload["message"]),
                #     self.loop
                # )
        except Exception as e:
            logger.error(f"处理Redis消息失败: {e}")
    @staticmethod
    async def send_personal_message(message: str, websocket: WebSocket):
        if websocket.client_state.CONNECTED:
            await websocket.send_text(message)
ws_manager = WebsocketManager()

